# import os
# import torch
# from ultralytics import YOLO
# import argparse

# def train_model(data_yaml, model_config, epochs=100, imgsz=640, batch_size=16, device='0'):
#     """
#     训练YOLOv8疲劳驾驶检测模型
    
#     Args:
#         data_yaml: 数据集配置文件路径
#         model_config: 模型配置文件路径
#         epochs: 训练轮数
#         imgsz: 输入图像尺寸
#         batch_size: 批次大小
#         device: 训练设备
#     """
    
#     # 设置设备
#     if device:
#         os.environ['CUDA_VISIBLE_DEVICES'] = device
    
#     # 加载模型
#     model = YOLO(model_config)
    
#     # 训练参数配置
#     train_args = {
#         'data': data_yaml,
#         'epochs': epochs,
#         'imgsz': imgsz,
#         'batch': batch_size,
#         'patience': 50,  # 早停耐心值
#         'save': True,
#         'save_period': 10,  # 每10个epoch保存一次
#         'cache': False,
#         'device': device,
#         'workers': 8,
#         'project': 'fatigue_detection',
#         'name': 'yolov8_fatigue',
#         'exist_ok': True,
#         'pretrained': True,
#         'optimizer': 'auto',
#         'lr0': 0.01,  # 初始学习率
#         'lrf': 0.01,  # 最终学习率
#         'momentum': 0.937,
#         'weight_decay': 0.0005,
#         'warmup_epochs': 3.0,
#         'warmup_momentum': 0.8,
#         'box': 7.5,
#         'cls': 0.5,
#         'dfl': 1.5,
#         'fl_gamma': 0.0,
#     }
    
#     # 开始训练
#     results = model.train(**train_args)
    
#     # 验证模型
#     val_results = model.val()
#     print(f"验证结果: {val_results}")
    
#     return model, results

# if __name__ == "__main__":
#     parser = argparse.ArgumentParser(description='YOLOv8疲劳驾驶检测训练')
#     parser.add_argument('--data', type=str, default='../dataset/data.yaml', help='数据集配置文件路径')
#     parser.add_argument('--model', type=str, default='../yolov8n.pt', help='模型配置文件路径')
#     parser.add_argument('--epochs', type=int, default=100, help='训练轮数')
#     parser.add_argument('--imgsz', type=int, default=640, help='输入图像尺寸')
#     parser.add_argument('--batch', type=int, default=16, help='批次大小')
#     parser.add_argument('--device', type=str, default='0', help='训练设备')
    
#     args = parser.parse_args()
    
#     # 开始训练
#     model, results = train_model(
#         data_yaml=args.data,
#         model_config=args.model,
#         epochs=args.epochs,
#         imgsz=args.imgsz,
#         batch_size=args.batch,
#         device=args.device
#     )
    
#     print("训练完成!")

import os
import torch
from ultralytics import YOLO
import argparse
import sys

def check_gpu_availability():
    """检查GPU可用性"""
    if torch.cuda.is_available():
        gpu_count = torch.cuda.device_count()
        print(f"🎯 检测到 {gpu_count} 个GPU:")
        for i in range(gpu_count):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
            print(f"  GPU {i}: {gpu_name} ({gpu_memory:.1f} GB)")
        return True
    else:
        print("❌ 未检测到GPU，将使用CPU进行训练")
        return False

def train_model(data_yaml, model_name='yolov8n.pt', epochs=100, imgsz=640, batch_size=16, device='auto'):
    """
    训练YOLOv8疲劳驾驶检测模型
    
    Args:
        data_yaml: 数据集配置文件路径
        model_name: 预训练模型名称或路径
        epochs: 训练轮数
        imgsz: 输入图像尺寸
        batch_size: 批次大小
        device: 训练设备 ('auto', 'cpu', or GPU ID)
    """
    
    # 自动检测设备
    if device == 'auto':
        if check_gpu_availability():
            device = '0'  # 使用第一个GPU
        else:
            device = 'cpu'
    
    # 设置设备
    if device != 'cpu':
        os.environ['CUDA_VISIBLE_DEVICES'] = device
    
    print(f"🚀 开始训练配置:")
    print(f"   模型: {model_name}")
    print(f"   数据集: {data_yaml}")
    print(f"   设备: {'GPU ' + device if device != 'cpu' else 'CPU'}")
    print(f"   训练轮数: {epochs}")
    print(f"   批次大小: {batch_size}")
    print(f"   图像尺寸: {imgsz}")
    
    try:
        # 加载预训练模型
        model = YOLO(model_name)
        print("✅ 模型加载成功!")
        
    except Exception as e:
        print(f"❌ 模型加载失败: {e}")
        print("🔄 尝试使用在线预训练模型...")
        try:
            model = YOLO('yolov8n.pt')
        except Exception as e2:
            print(f"❌ 在线模型也加载失败: {e2}")
            return None, None

    # 检查数据集文件是否存在
    if not os.path.exists(data_yaml):
        print(f"❌ 错误: 数据集配置文件不存在: {data_yaml}")
        print("请确保data.yaml文件存在")
        return None, None

    # 根据设备调整参数
    if device == 'cpu':
        batch_size = max(1, batch_size // 4)  # CPU上使用更小的批次
        workers = 2  # CPU上减少工作进程
        print(f"⚠️  使用CPU训练，调整批次大小为: {batch_size}")
    else:
        workers = 4 if os.cpu_count() > 4 else max(1, os.cpu_count() - 1)

    # 训练参数配置
    train_args = {
        'data': data_yaml,
        'epochs': epochs,
        'imgsz': imgsz,
        'batch': batch_size,
        'patience': 20,
        'save': True,
        'save_period': 10,
        'cache': False,
        'device': device,
        'workers': workers,
        'project': 'runs/detect',
        'name': 'fatigue_detection',
        'exist_ok': True,
        'pretrained': True,
        'optimizer': 'auto',
        'lr0': 0.01,
        'lrf': 0.01,
        'momentum': 0.937,
        'weight_decay': 0.0005,
        'warmup_epochs': 3.0,
        'warmup_momentum': 0.8,
        'box': 7.5,
        'cls': 0.5,
        'dfl': 1.5,
        'verbose': False,  # 减少输出
    }
    
    print("🔥 开始训练...")
    
    # 开始训练
    try:
        results = model.train(**train_args)
        
        if results is not None:
            # 验证模型
            print("📊 开始验证模型...")
            val_results = model.val()
            
            print("✅ 训练完成!")
            print("📈 验证结果:")
            print(f"   精确度 (mAP50): {val_results.box.map50:.4f}")
            print(f"   精确度 (mAP50-95): {val_results.box.map:.4f}")
            print(f"   准确率: {val_results.box.mp:.4f}")
            print(f"   召回率: {val_results.box.mr:.4f}")
            
            # 保存最佳模型到指定位置
            best_model_path = 'best_fatigue_detection.pt'
            try:
                # 复制最佳模型
                import shutil
                best_model_source = results.save_dir / 'weights' / 'best.pt'
                if best_model_source.exists():
                    shutil.copy2(best_model_source, best_model_path)
                    print(f"💾 最佳模型已保存到: {best_model_path}")
                else:
                    print("⚠️  未找到最佳模型文件")
            except Exception as e:
                print(f"⚠️  保存最佳模型时出错: {e}")
            
            return model, results
        else:
            print("❌ 训练结果为空")
            return None, None
        
    except Exception as e:
        print(f"❌ 训练过程中出现错误: {e}")
        import traceback
        traceback.print_exc()
        return None, None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='YOLOv8疲劳驾驶检测训练')
    parser.add_argument('--data', type=str, default='dataset/data.yaml', help='数据集配置文件路径')
    parser.add_argument('--model', type=str, default='yolov8n.pt', help='预训练模型名称或路径')
    parser.add_argument('--epochs', type=int, default=50, help='训练轮数')
    parser.add_argument('--imgsz', type=int, default=640, help='输入图像尺寸')
    parser.add_argument('--batch', type=int, default=8, help='批次大小')
    parser.add_argument('--device', type=str, default='auto', help='训练设备 (auto, cpu, or GPU ID)')
    
    args = parser.parse_args()
    
    # 开始训练
    model, results = train_model(
        data_yaml=args.data,
        model_name=args.model,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch_size=args.batch,
        device=args.device
    )
    
    if model is not None:
        print("🎉 训练完成!")
        print("📁 训练结果保存在: runs/detect/fatigue_detection/")
    else:
        print("💥 训练失败!")